import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.VoidFunction;

import java.io.FileInputStream;
import java.io.IOException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.FileHandler;
import java.util.logging.Logger;
import java.util.logging.SimpleFormatter;

public class HdfsSparkSyncWithValidationOptions {
    
    private final FileSystem srcFs;
    private final FileSystem dstFs;
    private final long bandwidthLimit;
    private static final Logger logger = Logger.getLogger("HdfsSparkSyncLog");

    public HdfsSparkSyncWithValidationOptions(Configuration srcConf, Configuration dstConf, long bandwidthLimit) throws IOException {
        srcFs = FileSystem.get(srcConf);
        dstFs = FileSystem.get(dstConf);
        this.bandwidthLimit = bandwidthLimit * 1024 * 1024; // 转换为字节/秒
        setupLogger();
    }

    private void setupLogger() throws IOException {
        FileHandler fh = new FileHandler("/tmp/HdfsSparkSync.log", true); // 临时日志文件
        fh.setFormatter(new SimpleFormatter());
        logger.addHandler(fh);
        logger.setUseParentHandlers(false);
    }

    /**
     * update 方法用于同步源目录与目标目录的文件差异
     */
    public void update(Path srcDir, Path dstDir, JavaSparkContext sc) throws IOException {
        List<Path> filesToCopy = new ArrayList<>();
        gatherFilesToUpdate(srcDir, dstDir, filesToCopy);
        
        JavaRDD<Path> filesRDD = sc.parallelize(filesToCopy);
        filesRDD.foreach((VoidFunction<Path>) path -> {
            Configuration srcConf = new Configuration();
            srcConf.set("hadoop.security.authentication", "Kerberos");
            FileSystem srcFs = FileSystem.get(srcConf);
            FileSystem dstFs = FileSystem.get(new Configuration());
            
            UserGroupInformation.setConfiguration(srcConf);
            UserGroupInformation.loginUserFromKeytab("your_principal@YOUR_REALM", "/path/to/your.keytab");

            Path dstPath = new Path(dstDir, path.getName());
            copyFile(srcFs, dstFs, path, dstPath);
            preservePermissions(srcFs, dstFs, path, dstPath);
        });
    }

    /**
     * delete 方法用于删除目标目录中不在源目录中的文件
     */
    public void delete(Path srcDir, Path dstDir) throws IOException {
        FileStatus[] dstFiles = dstFs.listStatus(dstDir);

        for (FileStatus dstFile : dstFiles) {
            Path dstPath = dstFile.getPath();
            Path srcPath = new Path(srcDir, dstPath.getName());

            if (dstFile.isDirectory()) {
                delete(srcPath, dstPath);
                if (dstFs.listStatus(dstPath).length == 0) {
                    dstFs.delete(dstPath, true);
                    logger.info("Deleted empty directory: " + dstPath);
                }
            } else {
                if (!srcFs.exists(srcPath)) {
                    dstFs.delete(dstPath, true);
                    logger.info("Deleted: " + dstPath);
                }
            }
        }
    }

    private void gatherFilesToUpdate(Path srcDir, Path dstDir, List<Path> filesToCopy) throws IOException {
        FileStatus[] srcFiles = srcFs.listStatus(srcDir);

        if (!dstFs.exists(dstDir)) {
            dstFs.mkdirs(dstDir);
        }

        for (FileStatus srcFile : srcFiles) {
            Path srcPath = srcFile.getPath();
            Path dstPath = new Path(dstDir, srcPath.getName());

            if (srcFile.isDirectory()) {
                gatherFilesToUpdate(srcPath, dstPath, filesToCopy);
            } else {
                if (!dstFs.exists(dstPath) || srcFs.getFileStatus(srcPath).getModificationTime() > dstFs.getFileStatus(dstPath).getModificationTime()) {
                    filesToCopy.add(srcPath);
                    logger.info("File to update: " + srcPath);
                }
            }
        }
    }

    private void copyFile(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        try (var inStream = srcFs.open(srcPath);
             var outStream = dstFs.create(dstPath, true)) {

            byte[] buffer = new byte[4096];
            long startTime = System.currentTimeMillis();
            int bytesRead;
            long totalBytesRead = 0;

            while ((bytesRead = inStream.read(buffer)) > 0) {
                outStream.write(buffer, 0, bytesRead);
                totalBytesRead += bytesRead;

                if (totalBytesRead >= bandwidthLimit) {
                    long elapsedTime = System.currentTimeMillis() - startTime;
                    if (elapsedTime < 1000) {
                        Thread.sleep(1000 - elapsedTime);
                    }
                    totalBytesRead = 0;
                    startTime = System.currentTimeMillis();
                }
            }
        }
        logger.info("Copied: " + srcPath + " to " + dstPath);
    }

    private void preservePermissions(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        FileStatus srcStatus = srcFs.getFileStatus(srcPath);
        dstFs.setPermission(dstPath, srcStatus.getPermission());
        dstFs.setOwner(dstPath, srcStatus.getOwner(), srcStatus.getGroup());
        logger.info("Permissions preserved for: " + dstPath);
    }

    /**
     * validateAndResyncFiles 方法，选择文件校验方式并重新同步不一致的文件
     */
    private void validateAndResyncFiles(List<Path> filesToCopy, Path dstDir, String validationMethod) throws IOException {
        for (Path srcPath : filesToCopy) {
            Path dstPath = new Path(dstDir, srcPath.getName());

            boolean isValid;
            switch (validationMethod.toLowerCase()) {
                case "md5":
                    isValid = validateChecksum(srcFs, dstFs, srcPath, dstPath);
                    break;
                case "size":
                    isValid = validateSize(srcFs, dstFs, srcPath, dstPath);
                    break;
                default:
                    throw new IllegalArgumentException("Unknown validation method: " + validationMethod);
            }

            if (!isValid) {
                logger.warning("Validation mismatch for: " + srcPath + ". Resyncing...");
                copyFile(srcFs, dstFs, srcPath, dstPath);
                preservePermissions(srcFs, dstFs, srcPath, dstPath);
            } else {
                logger.info("Validation successful for: " + srcPath);
            }
        }
    }

    private boolean validateSize(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        return srcFs.getFileStatus(srcPath).getLen() == dstFs.getFileStatus(dstPath).getLen();
    }

    private boolean validateChecksum(FileSystem srcFs, FileSystem dstFs, Path srcPath, Path dstPath) throws IOException {
        String srcChecksum = calculateChecksum(srcFs, srcPath);
        String dstChecksum = calculateChecksum(dstFs, dstPath);
        return srcChecksum.equals(dstChecksum);
    }

    private String calculateChecksum(FileSystem fs, Path path) throws IOException {
        try (FileInputStream fileStream = new FileInputStream(fs.open(path).getWrappedStream())) {
            MessageDigest md = MessageDigest.getInstance("MD5");
            byte[] buffer = new byte[4096];
            int bytesRead;

            while ((bytesRead = fileStream.read(buffer)) != -1) {
                md.update(buffer, 0, bytesRead);
            }
            StringBuilder checksum = new StringBuilder();
            for (byte b : md.digest()) {
                checksum.append(String.format("%02x", b));
            }
            return checksum.toString();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("MD5 algorithm not available", e);
        }
    }

    public void uploadLogToHdfs(Path hdfsLogDir) throws IOException {
        Path localLogPath = new Path("/tmp/HdfsSparkSync.log");
        Path hdfsLogPath = new Path(hdfsLogDir, "HdfsSparkSync.log");
        
        if (dstFs.exists(hdfsLogPath)) {
            dstFs.delete(hdfsLogPath, true);
        }
        
        dstFs.copyFromLocalFile(localLogPath, hdfsLogPath);
        System.out.println("Log file uploaded to HDFS: " + hdfsLogPath);
    }

    /**
     * main 方法：用于执行文件同步和校验，并选择校验方式
     */
    public static void main(String[] args) throws Exception {
        Configuration srcConf = new Configuration();
        Configuration dstConf = new Configuration();

        srcConf.set("fs.defaultFS", "hdfs://source_cluster");
        dstConf.set("fs.defaultFS", "hdfs://destination_cluster");

        UserGroupInformation.setConfiguration(srcConf);
        UserGroupInformation.loginUserFromKeytab("your_principal@YOUR_REALM", "/path/to/your.keytab");

        JavaSparkContext sc = new JavaSparkContext();

        HdfsSparkSyncWithValidationOptions syncTool = new HdfsSparkSyncWithValidationOptions(srcConf, dstConf, 20);

        Path srcDir = new Path("/source/path");
        Path dstDir = new Path("/destination/path");
        Path dfsLogDir = new Path("/destination/log/path");

        syncTool.update(srcDir, dstDir, sc);
        syncTool.delete(srcDir, dstDir);

        String validationMethod = "md5"; // 或使用 "size"
        syncTool.validateAndResyncFiles(List.of(srcDir), dstDir, validationMethod);

        syncTool.uploadLogToHdfs(dfsLogDir);

        sc.close();
    }
}
